/* *
 * Buddy System 内存分配算法 - 修复版本
 * */

#include <pmm.h>
#include <list.h>
#include <string.h>
#include <stdio.h>
#include <buddy_pmm.h>
#include <assert.h>

// 伙伴系统数据结构
struct buddy {
    unsigned int size;          // 管理的内存总大小（页数）
    unsigned int max_order;     // 最大阶数
    list_entry_t free_list[MAX_ORDER + 1];  // 各阶空闲链表
};

static struct buddy buddy_sys;

// 辅助函数：判断是否为2的幂
static bool is_power_of_2(unsigned int n) {
    return n != 0 && (n & (n - 1)) == 0;
}

// 辅助函数：向上取到2的幂
static unsigned int round_up_power_of_2(unsigned int n) {
    if (n == 0) return 1;
    if (is_power_of_2(n)) return n;
    
    unsigned int order = 0;
    unsigned int temp = n - 1;
    while (temp > 0) {
        temp >>= 1;
        order++;
    }
    return 1 << order;
}

// 计算阶数
static unsigned int get_order(unsigned int n) {
    if (n == 0) return 0;
    unsigned int order = 0;
    unsigned int temp = n;
    while (temp > 1) {
        temp >>= 1;
        order++;
    }
    return order;
}

// 初始化伙伴系统
static void buddy_init(void) {
    cprintf("buddy_init: starting\n");
    for (int i = 0; i <= MAX_ORDER; i++) {
        list_init(&buddy_sys.free_list[i]);
    }
    buddy_sys.size = 0;
    buddy_sys.max_order = 0;
    cprintf("buddy_init: completed\n");
}

// 初始化内存映射
static void buddy_init_memmap(struct Page *base, size_t n) {
    assert(n > 0);
    cprintf("buddy_init_memmap: base=%p, n=%lu\n", base, n);
    
    // 只初始化前64个页面用于测试
    size_t init_count = (n < 64) ? n : 64;
    for (struct Page *p = base; p != base + init_count; p++) {
        assert(PageReserved(p));
        p->flags = 0;
        p->property = 0;
        set_page_ref(p, 0);
        SetPageProperty(p);
    }
    
    // 使用固定的小内存块进行测试
    unsigned int total_pages = 64;
    buddy_sys.size = total_pages;
    buddy_sys.max_order = get_order(total_pages);
    
    // 将内存块加入到最高阶空闲链表
    base->property = total_pages;
    list_add(&buddy_sys.free_list[buddy_sys.max_order], &(base->page_link));
    
    cprintf("buddy system init: total_pages = %u, max_order = %u\n", 
            total_pages, buddy_sys.max_order);
    cprintf("buddy_init_memmap: completed successfully\n");
}

// 分配页面 - 修复边界条件
static struct Page *buddy_alloc_pages(size_t n) {
    // 检查边界条件
    if (n == 0 || n > buddy_sys.size) {
        return NULL;
    }
    
    // 计算需要的阶数
    unsigned int alloc_size = round_up_power_of_2(n);
    unsigned int order = get_order(alloc_size);
    
    // 如果当前阶没有空闲块，从更高阶分裂
    unsigned int current_order = order;
    while (current_order <= buddy_sys.max_order) {
        if (!list_empty(&buddy_sys.free_list[current_order])) {
            break;
        }
        current_order++;
    }
    
    if (current_order > buddy_sys.max_order) {
        return NULL; // 内存不足
    }
    
    // 从找到的阶分配
    list_entry_t *le = list_next(&buddy_sys.free_list[current_order]);
    struct Page *page = le2page(le, page_link);
    list_del(le);
    
    // 如果需要分裂，一直分裂到目标阶
    while (current_order > order) {
        current_order--;
        unsigned int split_size = 1 << current_order;
        
        // 分裂为两个伙伴块
        struct Page *buddy = page + split_size;
        buddy->property = split_size;
        
        // 将伙伴块加入到对应阶的空闲链表
        list_add(&buddy_sys.free_list[current_order], &(buddy->page_link));
    }
    
    page->property = alloc_size;
    ClearPageProperty(page); // 标记为已分配
    
    // 更新已分配页面计数
    for (struct Page *p = page; p != page + alloc_size; p++) {
        set_page_ref(p, 0);
    }
    
    return page;
}

// 检查伙伴块是否在空闲链表中
static int is_buddy_in_free_list(struct Page *buddy, unsigned int order) {
    if (buddy < (struct Page *)0 || buddy >= (struct Page *)0 + buddy_sys.size) {
        return 0;
    }
    
    if (!PageProperty(buddy) || buddy->property != (1 << order)) {
        return 0;
    }
    
    // 检查是否在对应的空闲链表中
    list_entry_t *le = &buddy_sys.free_list[order];
    list_entry_t *list_head = le;
    le = list_next(le);
    
    while (le != list_head) {
        struct Page *page = le2page(le, page_link);
        if (page == buddy) {
            return 1;
        }
        le = list_next(le);
    }
    
    return 0;
}

// 释放页面 - 修复边界条件
static void buddy_free_pages(struct Page *base, size_t n) {
    // 检查边界条件
    if (n == 0) {
        return;
    }
    
    // 使用分配时记录的大小，而不是参数n
    unsigned int free_size = base->property;
    
    // 标记页面为空闲
    for (struct Page *p = base; p != base + free_size; p++) {
        assert(!PageReserved(p));
        SetPageProperty(p);
        set_page_ref(p, 0);
    }
    
    unsigned int order = get_order(free_size);
    struct Page *current_block = base;
    unsigned int current_order = order;
    
    // 尝试合并伙伴块
    while (current_order < buddy_sys.max_order) {
        // 计算伙伴块地址
        unsigned long block_index = (current_block - (struct Page *)0);
        unsigned long buddy_index = block_index ^ (1 << current_order);
        struct Page *buddy = (struct Page *)0 + buddy_index;
        
        // 检查伙伴块是否可以合并
        if (!is_buddy_in_free_list(buddy, current_order)) {
            break;
        }
        
        // 从链表中移除伙伴块
        list_del(&(buddy->page_link));
        
        // 合并两个块（取地址较小的那个）
        if (current_block > buddy) {
            struct Page *temp = current_block;
            current_block = buddy;
        }
        
        current_order++;
    }
    
    // 将合并后的块加入链表
    current_block->property = 1 << current_order;
    list_add(&buddy_sys.free_list[current_order], &(current_block->page_link));
}

// 检查伙伴系统状态
static size_t buddy_nr_free_pages(void) {
    size_t total_free = 0;
    
    for (int i = 0; i <= buddy_sys.max_order; i++) {
        list_entry_t *le = &buddy_sys.free_list[i];
        list_entry_t *list_head = le;
        le = list_next(le);
        
        while (le != list_head) {
            struct Page *page = le2page(le, page_link);
            total_free += page->property;
            le = list_next(le);
        }
    }
    
    return total_free;
}

// 检查伙伴系统
static void buddy_check(void) {
    cprintf("=== Buddy System Check ===\n");
    cprintf("Total managed pages: %u\n", buddy_sys.size);
    cprintf("Max order: %u\n", buddy_sys.max_order);
    
    for (int i = 0; i <= buddy_sys.max_order; i++) {
        int count = 0;
        list_entry_t *le = &buddy_sys.free_list[i];
        list_entry_t *list_head = le;
        le = list_next(le);
        
        while (le != list_head) {
            count++;
            le = list_next(le);
        }
        cprintf("Order %d: %d free blocks (size=%d)\n", i, count, 1 << i);
    }
    
    cprintf("Total free pages: %lu\n", buddy_nr_free_pages());
    cprintf("==========================\n");
}

// 伙伴系统内存管理器
const struct pmm_manager buddy_pmm_manager = {
    .name = "buddy_pmm_manager",
    .init = buddy_init,
    .init_memmap = buddy_init_memmap,
    .alloc_pages = buddy_alloc_pages,
    .free_pages = buddy_free_pages,
    .nr_free_pages = buddy_nr_free_pages,
    .check = buddy_check,
};